from datetime import datetime
import asyncio
import logging
from dataclasses import dataclass, field
from typing import Type, Callable, Dict, List, Union, Awaitable, Any, Optional

from qtorch.events.bus import Event, EventBus

@dataclass
class EnhancedEvent(Event):
    """增强版基础事件类，支持更多元数据"""
    metadata: Dict[str, Any] = field(default_factory=dict)
    source: str = "unknown"
    
    def add_metadata(self, key: str, value: Any) -> None:
        """添加元数据"""
        self.metadata[key] = value
        
    def get_metadata(self, key: str, default: Any = None) -> Any:
        """获取元数据"""
        return self.metadata.get(key, default)

@dataclass
class MultiAssetMarketDataEvent(EnhancedEvent):
    """多资产市场数据事件"""
    symbols: Dict[str, str] = field(default_factory=dict)  # 符号到资产类型的映射
    data: Dict[str, Any] = field(default_factory=dict)      # 每个符号的数据
    asset_types: Dict[str, str] = field(default_factory=dict)  # 每个符号的资产类型
    
    def add_asset_data(self, symbol: str, data: Any, asset_type: str = "stock") -> None:
        """添加单个资产数据"""
        self.symbols[symbol] = asset_type
        self.data[symbol] = data
        self.asset_types[symbol] = asset_type
        
    def get_asset_data(self, symbol: str) -> Any:
        """获取单个资产数据"""
        return self.data.get(symbol)
        
    def get_asset_type(self, symbol: str) -> str:
        """获取资产类型"""
        return self.asset_types.get(symbol, "stock")
        
    def get_assets_by_type(self, asset_type: str) -> List[str]:
        """获取指定类型的所有资产"""
        return [symbol for symbol, atype in self.asset_types.items() if atype == asset_type]

@dataclass
class MultiAssetSignalEvent(EnhancedEvent):
    """多资产信号事件"""
    strategy: str = ""
    signals: Dict[str, Dict[str, Any]] = field(default_factory=dict)  # 每个符号的信号和详细信息
    
    def add_signal(self, symbol: str, signal_value: float, 
                   asset_type: str = "stock", details: Dict[str, Any] = None) -> None:
        """添加单个资产信号"""
        if details is None:
            details = {}
            
        self.signals[symbol] = {
            "value": signal_value,
            "asset_type": asset_type,
            **details
        }
        
    def get_signal(self, symbol: str) -> Dict[str, Any]:
        """获取单个资产信号"""
        return self.signals.get(symbol, {"value": 0})
        
    def get_signal_value(self, symbol: str) -> float:
        """获取信号值"""
        signal = self.signals.get(symbol, {"value": 0})
        return signal.get("value", 0)
        
    def get_signals_by_type(self, asset_type: str) -> Dict[str, Dict[str, Any]]:
        """获取指定类型的所有资产信号"""
        return {symbol: signal for symbol, signal in self.signals.items() 
                if signal.get("asset_type") == asset_type}

@dataclass
class AssetCorrelationEvent(EnhancedEvent):
    """资产相关性事件"""
    correlation_matrix: Dict[str, Dict[str, float]] = field(default_factory=dict)
    time_period: str = "daily"  # 'daily', 'weekly', 'monthly'
    lookback_periods: int = 20  # 回溯周期数
    
    def set_correlation(self, symbol1: str, symbol2: str, value: float) -> None:
        """设置两个资产间的相关性"""
        if symbol1 not in self.correlation_matrix:
            self.correlation_matrix[symbol1] = {}
        if symbol2 not in self.correlation_matrix:
            self.correlation_matrix[symbol2] = {}
            
        self.correlation_matrix[symbol1][symbol2] = value
        self.correlation_matrix[symbol2][symbol1] = value
        
    def get_correlation(self, symbol1: str, symbol2: str) -> float:
        """获取两个资产间的相关性"""
        if symbol1 in self.correlation_matrix and symbol2 in self.correlation_matrix[symbol1]:
            return self.correlation_matrix[symbol1][symbol2]
        return 0.0
        
    def get_all_correlations(self, symbol: str) -> Dict[str, float]:
        """获取一个资产与所有其他资产的相关性"""
        return self.correlation_matrix.get(symbol, {})
        
    def get_highly_correlated(self, symbol: str, threshold: float = 0.7) -> List[str]:
        """获取与指定资产高度相关的资产列表"""
        correlations = self.get_all_correlations(symbol)
        return [other_symbol for other_symbol, corr in correlations.items() 
                if corr >= threshold and other_symbol != symbol]

@dataclass
class PortfolioAllocationEvent(EnhancedEvent):
    """投资组合分配事件"""
    allocations: Dict[str, float] = field(default_factory=dict)  # 每个资产的权重分配
    risk_metrics: Dict[str, float] = field(default_factory=dict)  # 组合风险指标
    constraints: Dict[str, Any] = field(default_factory=dict)    # 分配约束条件
    
    def set_allocation(self, symbol: str, weight: float) -> None:
        """设置单个资产的权重"""
        self.allocations[symbol] = weight
        
    def get_allocation(self, symbol: str) -> float:
        """获取单个资产的权重"""
        return self.allocations.get(symbol, 0.0)
        
    def set_risk_metric(self, name: str, value: float) -> None:
        """设置风险指标"""
        self.risk_metrics[name] = value
        
    def get_risk_metric(self, name: str) -> float:
        """获取风险指标"""
        return self.risk_metrics.get(name, 0.0)
        
    def get_total_allocation(self) -> float:
        """获取总分配权重"""
        return sum(self.allocations.values())
        
    def normalize_allocations(self) -> None:
        """归一化权重使总和为1"""
        total = self.get_total_allocation()
        if total > 0:
            for symbol in self.allocations:
                self.allocations[symbol] /= total

@dataclass
class AssetInteractionEvent(EnhancedEvent):
    """资产交互事件，用于处理不同资产类型间的关系"""
    primary_symbol: str = ""
    primary_asset_type: str = "stock"
    related_assets: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    interaction_type: str = ""  # 例如：'hedge', 'spread', 'arbitrage'
    
    def add_related_asset(self, symbol: str, asset_type: str, relationship: str, 
                          details: Dict[str, Any] = None) -> None:
        """添加关联资产"""
        if details is None:
            details = {}
            
        self.related_assets[symbol] = {
            "asset_type": asset_type,
            "relationship": relationship,
            **details
        }
        
    def get_related_assets(self, relationship: Optional[str] = None) -> Dict[str, Dict[str, Any]]:
        """获取关联资产"""
        if relationship is None:
            return self.related_assets
        
        return {symbol: info for symbol, info in self.related_assets.items() 
                if info.get("relationship") == relationship}

class EnhancedEventBus(EventBus):
    """增强版事件总线，支持资产类型感知和多资产操作"""
    
    def __init__(self):
        super().__init__()
        self._asset_type_filters: Dict[Type[Event], Dict[str, List[Callable]]] = {}
        self._symbol_filters: Dict[Type[Event], Dict[str, List[Callable]]] = {}
        
    def subscribe_by_asset_type(self, event_type: Type[Event], asset_type: str, 
                                handler: Callable) -> None:
        """按资产类型订阅事件"""
        if event_type not in self._asset_type_filters:
            self._asset_type_filters[event_type] = {}
        
        if asset_type not in self._asset_type_filters[event_type]:
            self._asset_type_filters[event_type][asset_type] = []
            
        self._asset_type_filters[event_type][asset_type].append(handler)
        
    def subscribe_by_symbol(self, event_type: Type[Event], symbol: str, 
                           handler: Callable) -> None:
        """按标的符号订阅事件"""
        if event_type not in self._symbol_filters:
            self._symbol_filters[event_type] = {}
        
        if symbol not in self._symbol_filters[event_type]:
            self._symbol_filters[event_type][symbol] = []
            
        self._symbol_filters[event_type][symbol].append(handler)
        
    async def publish_async(self, event: Event) -> None:
        """异步发布事件"""
        await super().publish_async(event)
        
        # 处理资产类型特定的处理器
        if isinstance(event, (MultiAssetMarketDataEvent, MultiAssetSignalEvent)) and \
           type(event) in self._asset_type_filters:
            
            event_type = type(event)
            
            # 对于市场数据事件，按资产类型调用处理器
            if isinstance(event, MultiAssetMarketDataEvent):
                for symbol, asset_type in event.asset_types.items():
                    if asset_type in self._asset_type_filters[event_type]:
                        for handler in self._asset_type_filters[event_type][asset_type]:
                            handler(event)
            
            # 对于信号事件，按资产类型调用处理器
            elif isinstance(event, MultiAssetSignalEvent):
                for symbol, signal in event.signals.items():
                    asset_type = signal.get("asset_type", "stock")
                    if asset_type in self._asset_type_filters[event_type]:
                        for handler in self._asset_type_filters[event_type][asset_type]:
                            handler(event)
        
        # 处理符号特定的处理器
        if isinstance(event, (MultiAssetMarketDataEvent, MultiAssetSignalEvent)) and \
           type(event) in self._symbol_filters:
            
            event_type = type(event)
            
            # 对于市场数据事件，按符号调用处理器
            if isinstance(event, MultiAssetMarketDataEvent):
                for symbol in event.symbols:
                    if symbol in self._symbol_filters[event_type]:
                        for handler in self._symbol_filters[event_type][symbol]:
                            handler(event)
            
            # 对于信号事件，按符号调用处理器
            elif isinstance(event, MultiAssetSignalEvent):
                for symbol in event.signals:
                    if symbol in self._symbol_filters[event_type]:
                        for handler in self._symbol_filters[event_type][symbol]:
                            handler(event)